package top.wshape1.rpc.registry;

import lombok.extern.slf4j.Slf4j;
import org.apache.curator.framework.AuthInfo;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.framework.recipes.cache.ChildData;
import org.apache.curator.framework.recipes.cache.PathChildrenCache;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.apache.zookeeper.CreateMode;
import top.wshape1.rpc.common.entity.ServiceMeta;
import top.wshape1.rpc.common.constant.RoleType;
import top.wshape1.rpc.config.JJRpcConfig;
import top.wshape1.rpc.config.RegistryConfig;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

/**
 * @author Wshape1
 * @version 1.0
 * @description
 * @since 2024-02-04
 */

@Slf4j
public class ZookeeperRegistry implements RpcRegistry {

    /*
     * Constants
     */

    // 连接失败等待重试时间
    public static final int BASE_SLEEP_TIME = 1000;
    // 重试次数
    public static final int MAX_RETRIES = 3;

    /*
     * Zookeeper Curator Client
     */
    private final CuratorFramework client;

    /**
     * provider节点缓存
     * serviceName -> PathChildrenCache
     */
    private final ConcurrentHashMap<String, PathChildrenCache> providerCacheMap = new ConcurrentHashMap<>();

    /*
     * Constructor
     */

    public ZookeeperRegistry() {
        RegistryConfig.ZookeeperRegistryConfig ZK_CONFIG = JJRpcConfig.CONFIG.getRegistry().getZookeeper();
        try {
            CuratorFrameworkFactory.Builder builder = CuratorFrameworkFactory.builder()
                    .connectString(ZK_CONFIG.getAddress())
                    .sessionTimeoutMs(ZK_CONFIG.getSessionTimeout())
                    .connectionTimeoutMs(ZK_CONFIG.getConnectionTimeout())
                    .retryPolicy(new ExponentialBackoffRetry(BASE_SLEEP_TIME, MAX_RETRIES))
                    .namespace(ZK_CONFIG.getRootNode());
            // 添加acl
            List<String> aclList = ZK_CONFIG.getAcl();
            if (aclList != null && !aclList.isEmpty()) {
                ArrayList<AuthInfo> authInfos = new ArrayList<>();
                for (String acl : aclList) {
                    String[] split = acl.split(":");
                    if (split.length == 3)
                        authInfos.add(new AuthInfo(split[0], (split[1] + ":" + split[2]).getBytes()));
                }
                builder.authorization(authInfos);
            }
            this.client = builder.build();
            this.client.start();
        } catch (Exception e) {
            log.error("zookeeper registry init error");
            throw new RuntimeException(e.getMessage(), e);
        }
    }

    /*
     * Override Methods
     */

    @Override
    public void register(ServiceMeta serviceMeta) {
        try {
            this.client.create()
                    .creatingParentsIfNeeded()
                    .withMode(CreateMode.EPHEMERAL)
                    .forPath(toNodePath(serviceMeta));
            // 如果是注册消费者，订阅服务
            if (serviceMeta.getRole() == RoleType.consumer)
                subscribe(serviceMeta.getServiceName());
        } catch (Exception e) {
            log.error("Service {} '{}' register error: {}", serviceMeta.getRole().name(), serviceMeta.getServiceName(), e.getMessage(), e);
        }
    }

    @Override
    public void unregister(ServiceMeta serviceMeta) {
        try {
            this.client.delete()
                    .forPath(toNodePath(serviceMeta));
            // 如果是注销消费者，取消订阅
            if (serviceMeta.getRole() == RoleType.consumer)
                unsubscribe(serviceMeta.getServiceName());
        } catch (Exception e) {
            log.error("Service {} '{}' unregister error: {}", serviceMeta.getRole().name(), serviceMeta.getServiceName(), e.getMessage(), e);
        }
    }

    @Override
    public List<ServiceMeta> lookup(String serviceName) {
        PathChildrenCache cache = providerCacheMap.get(serviceName);
        if (cache != null) {
            List<ChildData> data = cache.getCurrentData();
            if (data != null && !data.isEmpty())
                return data.stream()
                        .map(childData -> {
                            String path = childData.getPath();
                            String address = path.substring(path.lastIndexOf("/") + 1);
                            ServiceMeta serviceMeta = new ServiceMeta();
                            serviceMeta.setAddress(address);
                            serviceMeta.setRole(RoleType.provider);
                            serviceMeta.setServiceName(serviceName);
                            return serviceMeta;
                        }).collect(Collectors.toList());
        }
        return new ArrayList<>(0);
    }

    @Override
    public void subscribe(String serviceName) {
        try {
            PathChildrenCache cache = providerCacheMap.get(serviceName);
            if (cache == null) {
                PathChildrenCache newCache = new PathChildrenCache(client, toProviderPath4Sub(serviceName), true);
//            newCache.getListenable().addListener();
                PathChildrenCache putRet = providerCacheMap.putIfAbsent(serviceName, newCache);
                if (putRet == null) {
                    // put成功
                    newCache.start(PathChildrenCache.StartMode.BUILD_INITIAL_CACHE);
                }
            }
        } catch (Exception e) {
            log.error("Subscribe provider '{}' failed. error: {}", serviceName, e.getMessage());
            throw new RuntimeException(e);
        }
    }

    @Override
    public void unsubscribe(String serviceName) {
        try {
            PathChildrenCache cache = providerCacheMap.remove(serviceName);
            if (cache != null) {
                cache.close();
            }
        } catch (Exception e) {
            log.error("Unsubscribe provider '{}' failed. error: {}", serviceName, e.getMessage());
            throw new RuntimeException(e);
        }
    }


    @Override
    public void destroy() {
        providerCacheMap.forEach((serviceName, cache) -> {
            try {
                cache.close();
            } catch (Exception e) {
                log.warn(e.getMessage(), e);
            }
        });
        providerCacheMap.clear();
        try {
            this.client.close();
        } catch (Exception e) {
            log.warn(e.getMessage(), e);
        }
    }

    /*
     * My Private Methods
     */

    private String toNodePath(ServiceMeta serviceMeta) {
        return "/" + serviceMeta.getServiceName() + "/"
                + serviceMeta.getRole().name() + "/"
                + serviceMeta.getAddress();
    }

    private String toProviderPath4Sub(String serviceName) {
        return "/" + serviceName + "/" + RoleType.provider.name();
    }


}
